import cv2

from random import randint
from PIL import Image
import numpy as np

import pycuda.driver as cuda
import pycuda.autoinit

import tensorrt as trt

import sys, os, glob
sys.path.insert(1, os.path.join(sys.path[0], ".."))
import common

# set up logging message level
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)

# Paths and file names go here, and to be referred later.
class ModelData(object):
    ENGINE_FILE = "engine"
    MODEL_FILE  = "lenet5.uff"
    INPUT_NAME  = "input_1"
    INPUT_SHAPE = (1, 28, 28)
    OUTPUT_NAME = "dense_1/Softmax"

def build_engine(model_file):


    # check if we have created a serialized engine.
    # If so, read and deserialize it. Otherwise,
    # create it now.
    if ('engine' in glob.glob('*')):
        print('Loading previously built engine...')
        with open('./engine', 'rb') as f, trt.Runtime(TRT_LOGGER) as runtime:
            return runtime.deserialize_cuda_engine(f.read())
    else:
        print('Creating a freshly new engine...')
        with trt.Builder(TRT_LOGGER) as builder, builder.create_network() as network, trt.UffParser() as parser:
            builder.max_workspace_size = common.GiB(1)

            parser.register_input(ModelData.INPUT_NAME, ModelData.INPUT_SHAPE)
            parser.register_output(ModelData.OUTPUT_NAME)
            parser.parse(model_file, network)

            engine = builder.build_cuda_engine(network)

            serialized = engine.serialize()

            with open('./engine', 'wb') as f:
                f.write(serialized)

            return engine



# Loads a test case into the provided pagelocked_buffer.
def load_normalized_test_case(data_path, pagelocked_buffer, case_num=randint(0, 9)):
    test_case_path = os.path.join(data_path, str(case_num) + ".pgm")
    img = np.array(Image.open(test_case_path)).ravel()
    np.copyto(pagelocked_buffer, 1.0 - img / 255.0)
    return test_case_path 

def main():
    data_path, _ = common.find_sample_data(description="Runs an MNIST network using a UFF model file", subfolder="mnist")
    model_path = os.environ.get("MODEL_PATH") or os.path.join(os.path.dirname(__file__), "models")
    model_file = os.path.join(model_path, ModelData.MODEL_FILE)

    with build_engine(model_file) as engine:

        inputs, outputs, bindings, stream = common.allocate_buffers(engine)

        with engine.create_execution_context() as context:
            test_case_path = load_normalized_test_case(data_path, pagelocked_buffer=inputs[0].host)
            [output] = common.do_inference(context, bindings=bindings, inputs=inputs, outputs=outputs, stream=stream)
            pred = np.argmax(output)

            image = cv2.imread(test_case_path)
            resized = cv2.resize(image, (0, 0), fx=10, fy=10)
            cv2.putText(resized, 'Prediction: ' + str(pred), (0, 25), cv2.FONT_HERSHEY_SIMPLEX ,  1, (0,0,0) , 2, cv2.LINE_AA)

            cv2.imshow('original image', resized)
            cv2.waitKey(0)

if __name__ == '__main__':
    main()
